#include "Utility.h"
#include <iostream>
#include <algorithm>
#include <unordered_map>
#include <unordered_set>
#include <set>
#include <vector>
using namespace std;

int QueryDist(vector<unordered_map<int,int> >& label, int node_a, int node_b){

	unordered_map<int,int>* label_left = &(label[node_a]);
	unordered_map<int,int>* label_right = &(label[node_b]);

	if(label[node_a].size() > label[node_b].size()){
		label_left =  &(label[node_b]);
		label_right = &(label[node_a]);
	}

	int best_dist = 1000000;

	for(auto it = label_left->begin(); it != label_left->end(); it++){
		int transit_node = it->first;
		auto right_it = label_right->find(transit_node);
		if(right_it == label_right->end()) continue;

		int dist = it->second + right_it->second;
		if(dist < best_dist) best_dist = dist;
	}

	return best_dist;
}

void PrunedBFS(int source_node,
	vector<vector<pair<int,int> > >& adj_lists, 
	vector<unordered_map<int,int> >& label,
	vector<int>& current_dist){

	unordered_set<int> visited_nodes; // to reset the current_dist vector to infinite for visited nodes
	unordered_set<int> confirmed_nodes; // node that are done

	set<pair<int,int> > distance_queue; // first element is distacne from source, second is the node index
	distance_queue.insert(make_pair(0,source_node));
	visited_nodes.insert(source_node);
	current_dist[source_node] = 0;

	unordered_map<int, unordered_map<int,int> > new_label;

	int num_ops = 0;
	while(!distance_queue.empty()){
		int best_node = distance_queue.begin()->second;
		int best_dist = distance_queue.begin()->first;

		distance_queue.erase(distance_queue.begin());
		
		/*cout << "best_node:" << best_node << endl;
		cout << "best_dist:" << best_dist << endl;*/

		if(confirmed_nodes.find(best_node) != confirmed_nodes.end()) continue; // the node has been visited before
		confirmed_nodes.insert(best_node);

		int query_dist = QueryDist(label,source_node,best_node);
		if(query_dist <= best_dist) continue; // no need to compute any more

		// append the label
		if(new_label.find(best_node) == new_label.end()){
			unordered_map<int,int> new_list;
			new_label[best_node] = new_list;
		}

		new_label[best_node].insert(make_pair(source_node,best_dist));

		for(int i = 0; i < adj_lists[best_node].size(); i++){
			int v = adj_lists[best_node][i].first;
			num_ops++;
			if(confirmed_nodes.find(v) != confirmed_nodes.end()) continue;
			visited_nodes.insert(v); // for reset current_dist

			int new_dist = current_dist[best_node] + adj_lists[best_node][i].second;
			if(new_dist < current_dist[v]){
				current_dist[v] = new_dist;
				distance_queue.insert(make_pair(new_dist,v));
			}
		}

	}

	// reset curret_dist
	for(auto it = visited_nodes.begin(); it != visited_nodes.end(); it++) current_dist[*it] = 1000000;

	// update the label
	for(auto it = new_label.begin(); it != new_label.end(); it++){
		int node_id = it->first;
		unordered_map<int,int>& label_list = it->second;

		for(auto dit = new_label[node_id].begin(); dit != new_label[node_id].end(); dit++){
			label[node_id].insert(*dit);
		}
	}
	cout << "number of ops: " << num_ops << " for node: " << source_node << " with degree: " << adj_lists[source_node].size() << endl;
}

int ShortestPath(int source, int target, vector<vector<pair<int,int> > >& adj_lists){
	
	set<pair<int,int> > distance_queue; 
	vector<int> current_dist(adj_lists.size(),1000000);
	unordered_set<int> determined_node;

	current_dist[source] = 0;
	distance_queue.insert(make_pair(0,source));
	
	while(!distance_queue.empty()){
		int best_node = distance_queue.begin()->second;
		int best_dist = distance_queue.begin()->first;

		if(best_node == target) return best_dist;

		distance_queue.erase(distance_queue.begin());
		if(determined_node.find(best_node) != determined_node.end()) continue;
		determined_node.insert(best_node);


		for(int i = 0; i < adj_lists[best_node].size(); i++){
			int v = adj_lists[best_node][i].first;
			if(determined_node.find(v) != determined_node.end()) continue;

			int alt = current_dist[best_node] + adj_lists[best_node][i].second;
			if(alt < current_dist[v]){
				current_dist[v] = alt;
				distance_queue.insert(make_pair(alt,v));
			}
		}
	}
	return -1;
}

int main(int argc, char *argv[]){

	if(argc != 2){
		cout << "wrong arguments for " << argv[0] << endl;
		cout << "1. input directory" << endl;
		exit(1);
	}

	string in_dir(argv[1]);

	vector<vector<pair<int,int> > > adj_lists;
	LoadAdjGraph(adj_lists,in_dir);

	//for(int i = 0; i < adj_lists.size(); i++){
	//	int node_a = i;
	//	cout << node_a << "$ ";
	//	for(int j = 0; j < adj_lists[i].size(); j++){
	//		int node_b = adj_lists[i][j].first;
	//		int dist = adj_lists[i][j].second;

	//		cout << node_b << ":" << dist << " ";
	//	}
	//	cout << endl;
	//}


	//for(int i = 0; i < adj_lists.size(); i++){
	//	int node_a = i;
	//	for(int j = 0; j < adj_lists[i].size(); j++){
	//		int node_b = adj_lists[i][j].first;
	//		int dist = adj_lists[i][j].second;

	//		vector<pair<int,int> >& target = adj_lists[node_b];
	//		bool found = false;
	//		for(int k = 0; k < target.size(); k++){
	//			if(target[k].first == node_a || target[k].second == dist){
	//				found = true;
	//			}
	//		}
	//		if(!found){
	//			cout << node_a << " " << node_b << " have problem" << endl;
	//			exit(1);
	//		}
	//	}
	//}

	vector<int> current_dist(adj_lists.size(),1000000);
	vector<unordered_map<int,int> > label(adj_lists.size());

	vector<pair<int,int> > node_rank;
	for(int i = 0; i < adj_lists.size(); i++) node_rank.push_back(make_pair(adj_lists[i].size(),i));
	sort(node_rank.begin(),node_rank.end(),pair_greater<int,int>());

	cout << "finish init: " << endl;
	for(int i = 0; i < node_rank.size(); i++){
		PrunedBFS(node_rank[i].second,adj_lists,label,current_dist);
	}

	//while(true){
	//	int node_a;
	//	int node_b;
	//	cin >> node_a;
	//	cin >> node_b;
	//	node_a--;
	//	node_b--;

	//	cout << "result from standard: "<< ShortestPath(node_a,node_b,adj_lists) << endl;
	//	cout << "result from fast: " << QueryDist(label,node_a,node_b) << endl;
	//}

	cout << "finish labeling" << endl;

	return 1;
}